"""
MyLangExtract 核心提取引擎

提供简化的信息提取接口
"""

from typing import List, Optional, Union, Dict, Any
import os

from .data import Extraction, ExampleData, ExtractResult
from ..providers.openai_compatible import (
    create_provider_from_preset, 
    create_custom_provider,
    OpenAICompatibleProvider
)


class Extractor:
    """信息提取器
    
    提供简化的信息提取接口，支持多种大模型服务商
    """
    
    def __init__(self, provider: OpenAICompatibleProvider):
        """初始化提取器
        
        Args:
            provider: 大模型提供商实例
        """
        self.provider = provider
    
    def extract(
        self,
        text: str,
        prompt_description: str,
        examples: List[ExampleData],
        use_function_calling: Optional[bool] = None
    ) -> ExtractResult:
        """执行信息提取
        
        Args:
            text: 要提取的文本
            prompt_description: 提取任务描述
            examples: 示例数据列表
            use_function_calling: 是否使用 Function Calling
            
        Returns:
            提取结果
        """
        if not examples:
            raise ValueError("必须提供至少一个示例")
        
        if not text.strip():
            return ExtractResult(
                text=text,
                extractions=[],
                success=True,
                metadata={"message": "输入文本为空"}
            )
        
        return self.provider.extract(
            text=text,
            prompt_description=prompt_description,
            examples=examples,
            use_function_calling=use_function_calling
        )
    
    def extract_batch(
        self,
        texts: List[str],
        prompt_description: str,
        examples: List[ExampleData],
        use_function_calling: Optional[bool] = None
    ) -> List[ExtractResult]:
        """批量提取信息
        
        Args:
            texts: 要提取的文本列表
            prompt_description: 提取任务描述
            examples: 示例数据列表
            use_function_calling: 是否使用 Function Calling
            
        Returns:
            提取结果列表
        """
        results = []
        for text in texts:
            result = self.extract(
                text=text,
                prompt_description=prompt_description,
                examples=examples,
                use_function_calling=use_function_calling
            )
            results.append(result)
        
        return results
    
    def get_provider_info(self) -> Dict[str, Any]:
        """获取提供商信息
        
        Returns:
            提供商信息字典
        """
        return {
            "model_name": self.provider.config.model_name,
            "base_url": self.provider.config.base_url,
            "max_tokens": self.provider.config.max_tokens,
            "temperature": self.provider.config.temperature,
            "supports_function_calling": self.provider.config.supports_function_calling
        }


def create_extractor(
    provider: Optional[str] = None,
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    model_name: Optional[str] = None,
    **kwargs
) -> Extractor:
    """创建提取器的便捷函数
    
    Args:
        provider: 预设提供商名称（如 "zhipu", "moonshot", "openai" 等）
        api_key: API 密钥
        base_url: 自定义 API 基础 URL
        model_name: 自定义模型名称
        **kwargs: 额外配置参数
        
    Returns:
        提取器实例
        
    Examples:
        # 使用预设提供商
        extractor = create_extractor(provider="zhipu", api_key="your_key")
        
        # 使用自定义配置
        extractor = create_extractor(
            base_url="https://api.custom.com/v1/",
            model_name="custom-model",
            api_key="your_key"
        )
        
        # 从环境变量获取配置
        extractor = create_extractor(provider="zhipu")  # 自动从 ZHIPU_API_KEY 获取
    """
    
    if provider:
        # 使用预设提供商
        provider_instance = create_provider_from_preset(
            provider_name=provider,
            api_key=api_key,
            **kwargs
        )
    elif base_url and model_name:
        # 使用自定义配置
        if not api_key:
            raise ValueError("使用自定义配置时必须提供 api_key")
        
        provider_instance = create_custom_provider(
            base_url=base_url,
            model_name=model_name,
            api_key=api_key,
            **kwargs
        )
    else:
        raise ValueError(
            "必须提供 provider 参数或者同时提供 base_url 和 model_name 参数"
        )
    
    return Extractor(provider_instance)


def quick_extract(
    text: str,
    prompt_description: str,
    examples: List[ExampleData],
    provider: str = "zhipu",
    api_key: Optional[str] = None,
    **kwargs
) -> ExtractResult:
    """快速提取函数，适用于一次性提取任务
    
    Args:
        text: 要提取的文本
        prompt_description: 提取任务描述
        examples: 示例数据列表
        provider: 提供商名称，默认为 "zhipu"
        api_key: API 密钥
        **kwargs: 额外配置参数
        
    Returns:
        提取结果
        
    Example:
        result = quick_extract(
            text="张三今年25岁，是一名软件工程师",
            prompt_description="提取人物信息",
            examples=[...],
            provider="zhipu"
        )
    """
    extractor = create_extractor(
        provider=provider,
        api_key=api_key,
        **kwargs
    )
    
    return extractor.extract(
        text=text,
        prompt_description=prompt_description,
        examples=examples
    )


# 便捷函数，用于创建示例数据
def create_example(text: str, *extractions: Extraction) -> ExampleData:
    """创建示例数据的便捷函数
    
    Args:
        text: 示例文本
        *extractions: 提取结果
        
    Returns:
        示例数据对象
        
    Example:
        example = create_example(
            "张三今年25岁，是一名软件工程师",
            Extraction("person", "张三", {"age": "25", "job": "软件工程师"})
        )
    """
    return ExampleData(text=text, extractions=list(extractions))


def create_extraction(
    extraction_class: str,
    extraction_text: str,
    **attributes
) -> Extraction:
    """创建提取结果的便捷函数
    
    Args:
        extraction_class: 提取类别
        extraction_text: 提取文本
        **attributes: 属性键值对
        
    Returns:
        提取结果对象
        
    Example:
        extraction = create_extraction(
            "person", "张三", 
            age="25", job="软件工程师"
        )
    """
    return Extraction(
        extraction_class=extraction_class,
        extraction_text=extraction_text,
        attributes=attributes if attributes else None
    )


# 用于调试和测试的函数
def list_available_providers() -> Dict[str, str]:
    """列出所有可用的预设提供商
    
    Returns:
        提供商名称到描述的映射
    """
    from ..providers.presets import list_available_providers
    return list_available_providers()


def validate_examples(examples: List[ExampleData]) -> List[str]:
    """验证示例数据的有效性
    
    Args:
        examples: 示例数据列表
        
    Returns:
        错误信息列表，如果为空则表示验证通过
    """
    errors = []
    
    if not examples:
        errors.append("示例列表不能为空")
        return errors
    
    for i, example in enumerate(examples):
        if not example.text.strip():
            errors.append(f"示例 {i+1} 的文本不能为空")
        
        if not example.extractions:
            errors.append(f"示例 {i+1} 必须包含至少一个提取结果")
        
        for j, extraction in enumerate(example.extractions):
            if not extraction.extraction_class.strip():
                errors.append(f"示例 {i+1} 提取结果 {j+1} 的类别不能为空")
            
            if not extraction.extraction_text.strip():
                errors.append(f"示例 {i+1} 提取结果 {j+1} 的文本不能为空")
            
            # 检查提取文本是否在原文中
            if extraction.extraction_text not in example.text:
                errors.append(
                    f"示例 {i+1} 提取结果 {j+1} 的文本 '{extraction.extraction_text}' "
                    f"不在原文中"
                )
    
    return errors


def test_provider_connection(
    provider: str,
    api_key: Optional[str] = None
) -> Dict[str, Any]:
    """测试提供商连接
    
    Args:
        provider: 提供商名称
        api_key: API 密钥
        
    Returns:
        测试结果字典
    """
    try:
        # 创建简单的测试示例
        test_example = create_example(
            "测试文本",
            create_extraction("test", "测试")
        )
        
        # 创建提取器
        extractor = create_extractor(provider=provider, api_key=api_key)
        
        # 执行简单测试
        result = extractor.extract(
            text="这是一个测试",
            prompt_description="提取测试信息",
            examples=[test_example]
        )
        
        return {
            "success": result.success,
            "provider_info": extractor.get_provider_info(),
            "error": result.error_message if not result.success else None
        }
        
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }
